from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware  # 解决跨域问题
import json
import random
import asyncio
from typing import Dict, List, Optional
import uuid

class ConnectionManager:
    def __init__(self):
        # 存储活跃连接：key=连接ID，value=(WebSocket对象, 最后心跳时间)
        self.active_connections: Dict[str, tuple[WebSocket, float]] = {}
        self.lock = asyncio.Lock()  # 保证并发安全

    async def connect(self, websocket: WebSocket) -> str:
        """建立连接，返回唯一连接ID"""
        await websocket.accept()
        conn_id = str(uuid.uuid4())  # 生成唯一ID
        async with self.lock:
            self.active_connections[conn_id] = (websocket, asyncio.get_event_loop().time())
        print(f"新客户端连接：{conn_id}，当前连接数：{len(self.active_connections)}")
        return conn_id

    def disconnect(self, conn_id: str):
        """断开连接"""
        asyncio.create_task(self._safe_disconnect(conn_id))

    async def _safe_disconnect(self, conn_id: str):
        async with self.lock:
            if conn_id in self.active_connections:
                del self.active_connections[conn_id]
                print(f"客户端断开：{conn_id}，当前连接数：{len(self.active_connections)}")

    async def send_message(self, conn_id: str, message: dict):
        """向单个客户端发送消息"""
        async with self.lock:
            if conn_id not in self.active_connections:
                return
            websocket, _ = self.active_connections[conn_id]
        try:
            await websocket.send_text(json.dumps(message))
        except Exception as e:
            print(f"向 {conn_id} 发送消息失败：{e}")
            self.disconnect(conn_id)

    async def broadcast(self, message: dict):
        """向所有活跃客户端广播消息"""
        async with self.lock:
            # 复制当前连接列表，避免迭代时修改
            connections = list(self.active_connections.values())
        for websocket, _ in connections:
            try:
                await websocket.send_text(json.dumps(message))
            except Exception as e:
                print(f"广播消息失败：{e}")

    async def update_heartbeat(self, conn_id: str):
        """更新客户端心跳时间"""
        async with self.lock:
            if conn_id in self.active_connections:
                self.active_connections[conn_id] = (
                    self.active_connections[conn_id][0],
                    asyncio.get_event_loop().time()
                )

    async def clean_expired_connections(self, timeout: float = 30.0):
        """清理超时未发心跳的连接（默认30秒）"""
        while True:
            await asyncio.sleep(10)  # 每10秒检查一次
            now = asyncio.get_event_loop().time()
            expired_ids = []
            async with self.lock:
                for conn_id, (_, last_heartbeat) in self.active_connections.items():
                    if now - last_heartbeat > timeout:
                        expired_ids.append(conn_id)
            # 断开超时连接
            for conn_id in expired_ids:
                print(f"客户端超时未发心跳，断开连接：{conn_id}")
                self.disconnect(conn_id)

# 2. 创建 FastAPI 应用实例
app = FastAPI(title="3D Trajectory Visualization Server")

# 初始化连接管理器
manager = ConnectionManager()

# 4. 模拟轨迹数据生成
#    这个函数会周期性地生成三维坐标点，并将其广播给所有客户端
async def generate_trajectory_data():
    """
    模拟一个沿螺旋线运动的点。
    在实际应用中，你可以替换为从传感器、文件或其他数据源读取数据。
    """
    t = 0  # 参数，用于计算螺旋线坐标
    while True:
        # 生成螺旋线的坐标
        x = t * 0.1
        y = 10 * asyncio.get_event_loop().time() % 5  # 模拟一些抖动
        z = 5 * (t % 10)

        # 准备要发送的数据，使用 JSON 格式
        trajectory_point = {
            "type": "circle_data",
            "timestamp": asyncio.get_event_loop().time(),
            "data": [{
                "x": x,
                "y": y,
                "z": z
            }]
        }
        
        # 将 Python 字典转换为 JSON 字符串
        message = json.dumps(trajectory_point)
        
        # 广播消息
        if manager.active_connections:  # 只有在有连接时才广播
            await manager.broadcast(trajectory_point)
            print(f"Broadcasted point: {message}")

        # 等待一段时间，控制数据发送频率（例如，每 100 毫秒发送一个点）
        await asyncio.sleep(0.2)
        t += 1

async def generate_trace_data():
    """模拟三维轨迹数据生成（3条轨迹，100ms/次推送）"""
    trace_ids = ["trace_1"]
    positions = {tid: [0.0, 0.0, 0.0] for tid in trace_ids}  # 初始位置

    while True:
        # 生成新位置（随机偏移，限制坐标范围）
        for tid in trace_ids:
            positions[tid][0] += random.uniform(-1.0, 1.0)
            positions[tid][1] += random.uniform(-2.0, 2.0)
            positions[tid][2] = 1
            # 限制在 [-15, 15] 范围内，避免轨迹过散
            positions[tid] = [max(-15, min(15, x)) for x in positions[tid]]

        # 构造消息（符合前端预期格式）
        message = {
            "type": "trace_data",
            "timestamp": asyncio.get_event_loop().time(),
            "data": [
                {"id": tid, "x": round(pos[0], 2), "y": round(pos[1], 2), "z": round(pos[2], 2)}
                for tid, pos in positions.items()
            ]
        }
        # print(f"Broadcasted point: {message}")
        # 广播给所有客户端
        await manager.broadcast(message)
        await asyncio.sleep(0.5)  # 100ms 推送一次

# -------------------------- WebSocket 端点 --------------------------
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # 1. 建立连接
    conn_id = await manager.connect(websocket)
    try:
        # 2. 循环接收客户端消息（心跳/控制指令）
        while True:
            data = await websocket.receive_text()
            try:
                msg = json.loads(data)
                # 处理心跳消息
                if msg.get("type") == "heartbeat":
                    await manager.update_heartbeat(conn_id)
                    # 回复心跳确认
                    await manager.send_message(conn_id, {"type": "heartbeat_ack"})
                # 处理前端控制指令（如重置轨迹）
                elif msg.get("type") == "reset_trace":
                    print(f"收到 {conn_id} 的重置轨迹指令")
                    # 可在这里添加后端轨迹重置逻辑（如清空历史数据）
                    await manager.send_message(conn_id, {"type": "reset_ack", "message": "轨迹已重置"})
                else:
                    print(f"收到未知消息：{data}")
            except json.JSONDecodeError:
                print(f"无效的JSON消息：{data}")
    except WebSocketDisconnect:
        # 3. 客户端主动断开
        manager.disconnect(conn_id)
    except Exception as e:
        # 4. 其他异常
        print(f"连接异常 {conn_id}：{e}")
        manager.disconnect(conn_id)

# -------------------------- 核心配置 --------------------------
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境可指定具体域名（如 "http://localhost:8000"）
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 静态文件服务（包含前端页面）
app.mount("/", StaticFiles(directory="static", html=True), name="static")

# -------------------------- 前端页面路由 --------------------------
@app.get("/")
async def get():
    return FileResponse("static/index.html")

# -------------------------- 服务启动/停止钩子 --------------------------
@app.on_event("startup")
async def startup_event():
    """服务启动时执行"""
    # 启动轨迹数据生成任务
    asyncio.create_task(generate_trace_data())

@app.on_event("shutdown")
async def shutdown_event():
    """服务停止时执行"""
    print("服务正在停止，断开所有客户端连接...")
    async with manager.lock:
        for conn_id in list(manager.active_connections.keys()):
            try:
                websocket, _ = manager.active_connections[conn_id]
                await websocket.close()
            except:
                pass
            del manager.active_connections[conn_id]

# 9. 主函数，用于通过 uvicorn 启动应用
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="192.168.10.124", port=8000, reload=True)